/*
Copyright 2025 PerfXLab (Beijing) Technologies Co., Ltd.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
#include <iostream>
#include <complex>
#include <vector>
#include <cmath>
#include <cstdlib> // for size_t
#include "test_common.hpp"

// ===================== 修复比较函数 =====================
// 浮点数比较
static inline bool nearly_equal(float a, float b, float eps = 1e-5f) {
    return std::fabs(a - b) < eps;
}

// 双精度浮点数比较
static inline bool nearly_equal(double a, double b, double eps = 1e-12) {
    return std::fabs(a - b) < eps;
}

// 复数类型比较（完全重写，避免模板链接错误）
static inline bool nearly_equal(std::complex<float> a, std::complex<float> b, float eps = 1e-5f) {
    return nearly_equal(a.real(), b.real(), eps) &&
           nearly_equal(a.imag(), b.imag(), eps);
}

static inline bool nearly_equal(std::complex<double> a, std::complex<double> b, double eps = 1e-12) {
    return nearly_equal(a.real(), b.real(), eps) &&
           nearly_equal(a.imag(), b.imag(), eps);
}

// ===================== 向量比较函数 =====================
// 重构为模板特化版本，避免推导失败
template <typename T>
bool compare_vectors(size_t n, const T* a, const T* b, 
                    decltype(nearly_equal(std::declval<T>(), std::declval<T>())) eps = 
                    std::is_same_v<T, float> ? 1e-5f : 
                    std::is_same_v<T, double> ? 1e-12 :
                    std::is_same_v<T, std::complex<float>> ? 1e-5f : 1e-12) 
{
    for (size_t i = 0; i < n; ++i) {
        if (!nearly_equal(a[i], b[i], eps)) {
            std::cout << "Mismatch at position " << i << ": "
                      << "expected " << b[i] << ", got " << a[i] << std::endl;
            return false;
        }
    }
    return true;
}

// ===================== 向量打印函数 =====================
template <typename T>
void print_vector(const char* name, size_t n, const T* v) {
    std::cout << name << ": [";
    for (size_t i = 0; i < n; ++i) {
        std::cout << v[i];
        if (i < n-1) std::cout << ", ";
    }
    std::cout << "]" << std::endl;
}

// ===================== 主测试函数 =====================
TEST(SparseBLASlevel2, gather) {
    int failures = 0;

    std::cout << "=== Testing OpenSpB_gather (Fully Fixed Version) ===" << std::endl;

    // ===================== float (sgather) 测试 =====================
    {
        const int nz = 3;
        float y[] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f};
        float x_actual[3] = {0};
        float x_expected[3] = {0};
        int indx[] = {0, 2, 4}; // 0-based索引

        // 调用目标函数
        OpenSpB_sgather(nz, y, 1, x_actual, indx, OPENSPBLAS_ZERO_BASE);
        
        // 调用参考函数
        BLAS_susga(nz, y, 1, x_expected, indx, blas_zero_base);
        
        if (!compare_vectors(static_cast<size_t>(nz), x_actual, x_expected)) {
            std::cout << "FAIL: sgather test" << std::endl;
            print_vector("Expected", static_cast<size_t>(nz), x_expected);
            print_vector("Actual", static_cast<size_t>(nz), x_actual);
            failures++;
        } else {
            std::cout << "PASS: sgather test" << std::endl;
        }
    }

    // ===================== double (dgather) 测试 =====================
    {
        const int nz = 4;
        double y[] = {1.1, 2.2, 3.3, 4.4, 5.5};
        double x_actual[4] = {0};
        double x_expected[4] = {0};
        int indx[] = {1, 2, 3, 4}; // 1-based索引

        // 调用目标函数
        OpenSpB_dgather(nz, y, 1, x_actual, indx, OPENSPBLAS_ONE_BASE);
        
        // 调用参考函数
        BLAS_dusga(nz, y, 1, x_expected, indx, blas_one_base);
        
        if (!compare_vectors(static_cast<size_t>(nz), x_actual, x_expected)) {
            std::cout << "FAIL: dgather test" << std::endl;
            failures++;
        } else {
            std::cout << "PASS: dgather test" << std::endl;
        }
    }

    // ===================== complex<float> (cgather) 测试 =====================
    {
        const int nz = 2;
        std::complex<float> y[] = {{1.0f, 2.0f}, {3.0f, 4.0f}, {5.0f, 6.0f}};
        std::complex<float> x_actual[2] = {0};
        std::complex<float> x_expected[2] = {0};
        int indx[] = {0, 2}; // 0-based索引

        // 调用目标函数
        OpenSpB_cgather(nz, y, 1, x_actual, indx, OPENSPBLAS_ZERO_BASE);
        
        // 调用参考函数
        BLAS_cusga(nz, y, 1, x_expected, indx, blas_zero_base);
        
        if (!compare_vectors(static_cast<size_t>(nz), x_actual, x_expected)) {
            std::cout << "FAIL: cgather test" << std::endl;
            print_vector("Expected", static_cast<size_t>(nz), x_expected);
            print_vector("Actual", static_cast<size_t>(nz), x_actual);
            failures++;
        } else {
            std::cout << "PASS: cgather test" << std::endl;
        }
    }

    // ===================== complex<double> (zgather) 测试 =====================
    {
        const int nz = 3;
        std::complex<double> y[] = {{1.1, 2.2}, {3.3, 4.4}, {5.5, 6.6}, {7.7, 8.8}};
        std::complex<double> x_actual[3] = {0};
        std::complex<double> x_expected[3] = {0};
        int indx[] = {1, 2, 3}; // 0-based索引，步长为2
        const int incy = 2;

        // 调用目标函数
        OpenSpB_zgather(nz, y, incy, x_actual, indx, OPENSPBLAS_ZERO_BASE);
        
        // 调用参考函数
        BLAS_zusga(nz, y, incy, x_expected, indx, blas_zero_base);
        
        if (!compare_vectors(static_cast<size_t>(nz), x_actual, x_expected)) {
            std::cout << "FAIL: zgather test" << std::endl;
            failures++;
        } else {
            std::cout << "PASS: zgather test" << std::endl;
        }
    }

    // ===================== 空输入测试（修复零长度数组问题） =====================
    {
        float y[] = {1.0f, 2.0f, 3.0f};
        float x_actual;
        
        // 修复零长度数组问题 - 使用指针代替
        int* indx = nullptr;
        
        // 应安全处理空输入
        OpenSpB_sgather(0, y, 1, &x_actual, indx, OPENSPBLAS_ZERO_BASE);
        std::cout << "PASS: empty input test (no crash)" << std::endl;
    }

    // ===================== 测试总结 =====================
    std::cout << "\n=== Test Summary ===" << std::endl;
    if (failures == 0) {
        std::cout << "ALL TESTS PASSED!" << std::endl;
    } else {
        std::cout << failures << " TESTS FAILED!" << std::endl;
    }
}